# src/utils/data_validator.py

from datetime import datetime, timezone
from typing import List, Dict
import pandas as pd

from src.data_access.base_repository import BaseRepository
from src.data_access.timescale_repository import TimescaleRepository
from src.enums.timeframe import Timeframe


class DataValidator:
    """
    一个用于校验K线数据完整性的工具。
    它可以检测在给定时间范围内是否存在时间戳缺失。
    """
    # 将我们的Timeframe枚举映射到Pandas能理解的频率字符串
    TIMEFRAME_TO_FREQ_MAP = {
        Timeframe.TIMEFRAME_1M: "1min",
        Timeframe.TIMEFRAME_3M: "3min",
        Timeframe.TIMEFRAME_5M: "5min",
        Timeframe.TIMEFRAME_15M: "15min",
        Timeframe.TIMEFRAME_30M: "30min",
        Timeframe.TIMEFRAME_1H: "1h",
        Timeframe.TIMEFRAME_2H: "2h",
        Timeframe.TIMEFRAME_4H: "4h",
        Timeframe.TIMEFRAME_6H: "6h",
        Timeframe.TIMEFRAME_8H: "8h",
        Timeframe.TIMEFRAME_12H: "12h",
        Timeframe.TIMEFRAME_1D: "1D",
        Timeframe.TIMEFRAME_3D: "3D",
        Timeframe.TIMEFRAME_1W: "1W",
        # 注意: '月' 不是固定频率，校验月度数据意义不大，故在此省略
    }

    def __init__(self, repository: BaseRepository):
        """
        初始化校验器。

        Args:
            repository: 用于从数据库获取数据的仓储实例。
        """
        self.repo = repository

    def _get_expected_timestamps(self, start_date: datetime, end_date: datetime,
                                 timeframe: Timeframe) -> pd.DatetimeIndex:
        """生成理论上完整的、完美的日期时间索引。"""
        freq = self.TIMEFRAME_TO_FREQ_MAP.get(timeframe)
        if not freq:
            raise ValueError(f"不支持的时间周期校验: {timeframe.value}")

        # closed='left' 意味着包含起始点，不包含结束点，这与K线开盘时间的定义一致
        return pd.date_range(start=start_date, end=end_date, freq=freq, inclusive='left', tz='UTC')

    def validate_symbol(self, symbol: str, timeframes: List[Timeframe], start_date: datetime, end_date: datetime) -> \
            Dict[str, pd.DatetimeIndex]:
        """
        对单个交易对的多个时间周期进行数据完整性校验。

        Args:
            symbol: 要校验的交易对, e.g., 'BTC/USDT'。
            timeframes: 要校验的时间周期列表 (Timeframe枚举成员)。
            start_date: 校验的开始时间。
            end_date: 校验的结束时间。

        Returns:
            一个字典，键为时间周期字符串，值为缺失时间戳的DatetimeIndex。
        """
        print(f"--- 开始校验 {symbol} 从 {start_date.date()} 到 {end_date.date()} ---")

        report: Dict[str, pd.DatetimeIndex] = {}

        for tf in timeframes:
            print(f"  正在处理周期: {tf.value}...")

            # 1. 生成期望的时间戳
            expected_ts = self._get_expected_timestamps(start_date, end_date, tf)

            # 2. 获取实际存在的时间戳
            actual_ts = self.repo.get_kline_timestamps(symbol, tf.value, start_date, end_date)

            # 3. 找出差异 (缺失的数据)
            missing_ts = expected_ts.difference(actual_ts)

            report[tf.value] = missing_ts

            if missing_ts.empty:
                print(f"  ✅ {tf.value}: 数据完整，无缺失。")
            else:
                print(f"  ❌ {tf.value}: 发现 {len(missing_ts)} 个缺失数据点！")

        print(f"--- {symbol} 校验完成 ---\n")
        return report


if __name__ == '__main__':
    # --- 配置 ---
    DB_URL = "postgresql://postgres:postgres@127.0.0.1:15432/crypto?client_encoding=utf8"

    # 要校验的参数
    SYMBOL_TO_VALIDATE = "XRPUSDT"
    TIMEFRAMES_TO_VALIDATE = [
        Timeframe.TIMEFRAME_15M,
        Timeframe.TIMEFRAME_1H,
        Timeframe.TIMEFRAME_4H,
        Timeframe.TIMEFRAME_5M,
        Timeframe.TIMEFRAME_1M,
        Timeframe.TIMEFRAME_30M,
        Timeframe.TIMEFRAME_1D,
    ]

    # 使用带时区的时间是最佳实践
    START_DATE = datetime(2020, 1, 1, tzinfo=timezone.utc)
    END_DATE = datetime(2025, 8, 31, tzinfo=timezone.utc)


    def print_report(report: dict):
        """格式化并打印校验报告。"""
        has_missing_data = False
        print("=" * 50)
        print("      数 据 完 整 性 校 验 报 告")
        print("=" * 50)

        for timeframe, missing_ts in report.items():
            if not missing_ts.empty:
                has_missing_data = True
                print(f"\n❌ 在周期 [{timeframe}] 中发现缺失数据点:")
                # 为了防止打印过多，只显示前10个和后10个缺失点
                if len(missing_ts) > 20:
                    print(" (仅显示部分缺失点) ".center(40, '-'))
                    for ts in missing_ts[:10]:
                        print(f"  - {ts.strftime('%Y-%m-%d %H:%M:%S %Z')}")
                    print("  ...")
                    for ts in missing_ts[-10:]:
                        print(f"  - {ts.strftime('%Y-%m-%d %H:%M:%S %Z')}")
                else:
                    for ts in missing_ts:
                        print(f"  - {ts.strftime('%Y-%m-%d %H:%M:%S %Z')}")

        if not has_missing_data:
            print("\n✅ 恭喜！在指定范围内所有选定周期的数据都是完整的！")

        print("\n" + "=" * 50)


    # 1. 初始化 Repository
    try:
        repo = TimescaleRepository(db_url=DB_URL)
    except Exception:
        print("无法连接到数据库，请检查 DB_URL 和数据库服务。")
        raise

        # 2. 初始化 Validator
    validator = DataValidator(repository=repo)

    # 3. 运行校验任务
    validation_report = validator.validate_symbol(
        symbol=SYMBOL_TO_VALIDATE,
        timeframes=TIMEFRAMES_TO_VALIDATE,
        start_date=START_DATE,
        end_date=END_DATE
    )

    # 4. 打印格式化的报告
    print_report(validation_report)
